module ModuleElasticities
    
    use Globals
    use Functions_wa
    Use Functions_wa_TR
    
contains
    
    subroutine elasticities
    
        implicit none
        
        ! Local variable declarations
        real(8), dimension(NS) :: yk, yl, ytot, elasticity_num, elasticity_den, elasticity
        real(8), dimension(NS) :: ytot_sorted, mu_sorted, mu_cumu, elasticity_sorted, yl_sorted
        integer :: ind_sort(NS)
        real(8) :: elasticity_ytot_q(7), elasticity_yl_q(7)
        integer :: p_20_ind, p_40_ind, p_60_ind, p_80_ind, p_90_ind, p_99_ind
        integer :: is



        ! Save labor policies
        h_pol_epsilon(:,1) = h_pol

        ! Income
        yk = r*S(:,1)           ! capital income
        yl = wge*S(:,2)*h_pol   ! labor income
        ytot = yk+yl            ! total income

        ! taxes and transfers
        do is = 1,NS
            marginal_rate(is,1) = dTAX_sc(yl(is),yk(is))    ! marginal tax rate
        enddo
        
        ! store income
        ytot_epsilon(:,1) = ytot    ! total income
        yl_epsilon(:,1) = yl        ! labor income


        ! Period of the shock

        h_pol_epsilon(:,2) = h_pol_TR(:,1)

        ! Income
        yk = r_TR(1)*S(:,1)           ! capital income
        yl = wge_TR(1)*S(:,2)*h_pol_TR(:,1)   ! labor income
        ytot = yk+yl            ! total income

        ! taxes and transfers
        do is = 1,NS
            marginal_rate(is,2) = dTAX_sc_TR(yl(is),yk(is),1)    ! marginal tax rate
        enddo

        ! store income
        ytot_epsilon(:,2) = ytot    ! total income
        yl_epsilon(:,2) = yl        ! labor income



        ! Compute elasticities if index=2 [index=1: baseline]

        elasticity_num = (h_pol_epsilon(:,2)-h_pol_epsilon(:,1))/h_pol_epsilon(:,1)                     ! numerator
        elasticity_den = ((1d0-marginal_rate(:,2))-(1d0-marginal_rate(:,1)))/(1d0-marginal_rate(:,1))   ! denominator
        elasticity = elasticity_num/elasticity_den                                                      ! elasticity

        print *, 'elasticity_num [want negative number], max, mean, min = ', maxval(elasticity_num), sum(mu*elasticity_num), minval(elasticity_num)
        print *, 'elasticity_den [want negative number], max = ', maxval(elasticity_den)
        print *, 'elasticity [want positive], mean = ', sum(mu*elasticity)

        ! Sort by income, total
        ytot_sorted = ytot_epsilon(:,1)
        ind_sort = 0
        call hpsort_eps_epw(NS, ytot_sorted, ind_sort, 1D-16)
        mu_sorted = mu(ind_sort)
        elasticity_sorted = elasticity(ind_sort)

        ! compute cumulative measure
        mu_cumu = 0d0
        mu_cumu(1) = mu_sorted(1)
        do is = 2, NS
            mu_cumu(is) = mu_cumu(is-1) + mu_sorted(is)
        end do
        
        ! Percentile indices
        p_20_ind = minloc(abs(mu_cumu-0.2d0), dim=1)
        p_40_ind = minloc(abs(mu_cumu-0.4d0), dim=1)
        p_60_ind = minloc(abs(mu_cumu-0.6d0), dim=1)
        p_80_ind = minloc(abs(mu_cumu-0.8d0), dim=1)
        p_90_ind = minloc(abs(mu_cumu-0.9d0), dim=1)
        p_99_ind = minloc(abs(mu_cumu-0.99d0), dim=1)

        ! compute average elasticities by quintile
        elasticity_ytot_q(1) = sum(elasticity_sorted(1:p_20_ind)*mu_sorted(1:p_20_ind))/sum(mu_sorted(1:p_20_ind))
        elasticity_ytot_q(2) = sum(elasticity_sorted(p_20_ind+1:p_40_ind)*mu_sorted(p_20_ind+1:p_40_ind))/sum(mu_sorted(p_20_ind+1:p_40_ind))
        elasticity_ytot_q(3) = sum(elasticity_sorted(p_40_ind+1:p_60_ind)*mu_sorted(p_40_ind+1:p_60_ind))/sum(mu_sorted(p_40_ind+1:p_60_ind))
        elasticity_ytot_q(4) = sum(elasticity_sorted(p_60_ind+1:p_80_ind)*mu_sorted(p_60_ind+1:p_80_ind))/sum(mu_sorted(p_60_ind+1:p_80_ind))
        elasticity_ytot_q(5) = sum(elasticity_sorted(p_80_ind+1:NS)*mu_sorted(p_80_ind+1:NS))/sum(mu_sorted(p_80_ind+1:NS))
        elasticity_ytot_q(6) = sum(elasticity_sorted(p_90_ind:NS)*mu_sorted(p_90_ind:NS))/sum(mu_sorted(p_90_ind:NS))
        elasticity_ytot_q(7) = sum(elasticity_sorted(p_99_ind:NS)*mu_sorted(p_99_ind:NS))/sum(mu_sorted(p_99_ind:NS))

        ! Sort by income, labor
        yl_sorted = yl_epsilon(:,1)
        ind_sort = 0
        call hpsort_eps_epw(NS, yl_sorted, ind_sort, 1D-16)
        mu_sorted = mu(ind_sort)
        elasticity_sorted = elasticity(ind_sort)

        ! compute cumulative measure
        mu_cumu = 0d0
        mu_cumu(1) = mu_sorted(1)
        do is = 2, NS
         mu_cumu(is) = mu_cumu(is-1) + mu_sorted(is)
        end do

        ! percentile indices
        p_20_ind = minloc(abs(mu_cumu-0.2d0), dim=1)
        p_40_ind = minloc(abs(mu_cumu-0.4d0), dim=1)
        p_60_ind = minloc(abs(mu_cumu-0.6d0), dim=1)
        p_80_ind = minloc(abs(mu_cumu-0.8d0), dim=1)
        p_90_ind = minloc(abs(mu_cumu-0.9d0), dim=1)
        p_99_ind = minloc(abs(mu_cumu-0.99d0), dim=1)

        ! compute average elasticities by quintile
        elasticity_yl_q(1) = sum(elasticity_sorted(1:p_20_ind)*mu_sorted(1:p_20_ind))/sum(mu_sorted(1:p_20_ind))
        elasticity_yl_q(2) = sum(elasticity_sorted(p_20_ind+1:p_40_ind)*mu_sorted(p_20_ind+1:p_40_ind))/sum(mu_sorted(p_20_ind+1:p_40_ind))
        elasticity_yl_q(3) = sum(elasticity_sorted(p_40_ind+1:p_60_ind)*mu_sorted(p_40_ind+1:p_60_ind))/sum(mu_sorted(p_40_ind+1:p_60_ind))
        elasticity_yl_q(4) = sum(elasticity_sorted(p_60_ind+1:p_80_ind)*mu_sorted(p_60_ind+1:p_80_ind))/sum(mu_sorted(p_60_ind+1:p_80_ind))
        elasticity_yl_q(5) = sum(elasticity_sorted(p_80_ind+1:NS)*mu_sorted(p_80_ind+1:NS))/sum(mu_sorted(p_80_ind+1:NS))
        elasticity_yl_q(6) = sum(elasticity_sorted(p_90_ind:NS)*mu_sorted(p_90_ind:NS))/sum(mu_sorted(p_90_ind:NS))
        elasticity_yl_q(7) = sum(elasticity_sorted(p_99_ind:NS)*mu_sorted(p_99_ind:NS))/sum(mu_sorted(p_99_ind:NS))

        print *, 'elasticities by total income = ', elasticity_ytot_q
        print *, 'elasticities by labor income = ', elasticity_yl_q
    
    end subroutine elasticities
    
end module ModuleElasticities
    
